package org.example.common.threadpool;

import lombok.extern.slf4j.Slf4j;
import org.apache.skywalking.apm.toolkit.trace.RunnableWrapper;
import org.slf4j.MDC;
import org.springframework.core.task.TaskDecorator;
import org.springframework.web.context.request.RequestAttributes;
import org.springframework.web.context.request.RequestContextHolder;

import java.util.Map;

/**
 * 请求上下文包装类
 *
 * @author dingcan
 * @since 2021-04-29
 */
@Slf4j
public class ContextCopyingDecorator implements TaskDecorator {
    @Override
    public Runnable decorate(Runnable runnable) {
        try {
            RunnableWrapper runnableWithSW = RunnableWrapper.of(runnable);
            RequestAttributes context = RequestContextHolder.currentRequestAttributes();
            Map<String, String> previous = MDC.getCopyOfContextMap();
            return () -> {
                try {
                    RequestContextHolder.setRequestAttributes(context);
                    if (previous != null && !previous.isEmpty()) {
                        MDC.setContextMap(previous);
                    }
                    runnableWithSW.run();
                } finally {
                    RequestContextHolder.resetRequestAttributes();
                    MDC.clear();
                }
            };
        } catch (IllegalStateException e) {
            log.warn("ContextCopyingDecorator occur error:", e);
            return runnable;
        }
    }
}
